package com.ds.admin.config.thread;

import com.ds.common.trace.TraceIdContext;
import org.apache.commons.lang3.StringUtils;
import org.springframework.core.task.TaskDecorator;

/**
 * @author: fs
 * @date: 2023/7/20 13:01
 */
public class TraceIdTaskDecorator implements TaskDecorator {

	public Runnable decorator(Runnable runnable) {
		String traceId = TraceIdContext.getTraceId();
		return () -> {
			if (StringUtils.isNotBlank(traceId)) {
				try {
					TraceIdContext.setTraceId(traceId);
					runnable.run();
				}
				finally {
					TraceIdContext.clearTraceId();
				}
			}
			else {
				runnable.run();
			}
		};
	}

	@Override
	public Runnable decorate(Runnable runnable) {
		return decorator(runnable);
	}

}
